from fastapi import APIRouter
from pathlib import Path
import shutil,os,sys,fitz
import zipfile
from datetime import datetime
import random
import json
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.responses import JSONResponse
from concurrent.futures import ThreadPoolExecutor
from cnocr import CnOcr


requestPdfRename = APIRouter()
UPLOAD_DIR = Path("./upload")
JSON_FILE = UPLOAD_DIR / "translated.json"
UPLOAD_DIR.mkdir(exist_ok=True)
executor = ThreadPoolExecutor(max_workers=4)   

def pyMuPDF_fitz(pdfPath,imagePath):#pdf转成图片
    pdfDoc = fitz.open(pdfPath)
    page = pdfDoc[0]#只转第一页
    rotate = int(0)
    zoom_x = 1.3
    zoom_y = 1.3
    mat = fitz.Matrix(zoom_x, zoom_y).prerotate(rotate)
    pix = page.get_pixmap(matrix=mat, alpha=False)          
    pix.save (imagePath)   # 保存图片


def pdfRename(pdfPath,pdfName,res,stuNames,stuNames_dict,stuCard_dict):#重命名  
    pdfPath=str(pdfPath)+r"/"
    name=""#记录识别到的学生姓名
    pdfNewName = ""#记录目标文件名
    record=False#当遇到姓名或者身份证号时，将此变量设置为true，表示下一个内容要记录
    exitFile=False#记录完身份证后，将此变量设置为true，表示获取完成，可以退出当前文件的处理
    for i in range(len(res)):#遍历图片识别结果，是一个由字典组成的列表
        for k,v in res[i].items():#遍历识别出的每一项数据，数据是一个字典
            if k=="text":
                if record:
                    if exitFile:
                        cardid=str(v).strip()#去除识别出来的身份证前后的空格
                        if len(cardid)==0:#如果身份证为空，识别下一个
                            continue
                        pdfNewName=cardid+name
                        if pdfNewName in stuNames:#如果文件名在name.txt文件中，直接修改
                            os.rename(pdfPath + pdfName,pdfPath + pdfNewName + ".pdf")
                        elif str(v) in stuNames_dict.keys():#如果文件名不在，但是身份证号码正确，则修改姓名，以身份证号码+修改后的姓名作为文件名重命名
                            os.rename(pdfPath + pdfName,pdfPath+str(v)+stuNames_dict[str(v)]+".pdf")
                            shutil.copy(pdfPath+str(v) + stuNames_dict[str(v)] + ".pdf",pdfPath + "待检查/" + "姓名识别失败/")
                        elif name in stuCard_dict:
                            os.rename(pdfPath + pdfName,pdfPath + stuCard_dict[name] + name + ".pdf")
                            shutil.copy(pdfPath + stuCard_dict[name] + name + ".pdf",pdfPath + "待检查/" + "身份证识别失败/")
                        else:#如果身份证号码和姓名都不正确，不修改，提示手动修改
                            print("{}识别失败，识别为{},自动修改失败，请手动修改。".format(pdfName,pdfNewName))
                        return
                    name=str(v)#记录姓名
                    record=False
                elif v=="姓名":
                    record=True
                elif v=="身份证号":
                    record=True
                    exitFile=True
                else:
                    record=False

def _background_repack(zip_bytes: bytes, schoolName: str,studentInfo:str, code: str,tmp_zip:str) -> None:
    """线程里干苦力"""
    work_dir = UPLOAD_DIR / code

    with zipfile.ZipFile(tmp_zip, "r") as zf:
        zf.extractall(work_dir)
    tmp_zip.unlink()


    if os.path.exists(work_dir / "待检查"):
        shutil.rmtree(work_dir / "待检查")#删除待检查文件夹
    filenames=os.listdir(work_dir)#遍历存放pdf的文件夹，获取所有pdf文件名
    
    temp_dir = work_dir / "待检查"
    os.makedirs(temp_dir)
    os.makedirs(temp_dir / "姓名识别失败")
    os.makedirs(temp_dir / "身份证识别失败")

    studentInfo=studentInfo.split("\r\n")

    stuNames_dict={}
    stuCard_dict={} 
    for n in studentInfo:#建立身份信息字典，key为身份证，value为姓名，用于校对和纠正识别的结果
        stuNames_dict[n[0:18]]=n[18:]
    for k,v in stuNames_dict.items():#建立身份信息字典，key为姓名，value为身份证，用于校对和纠正识别的结果
        if v not in stuCard_dict.keys():#如果姓名在字典中不存在，则加入字典，如果已经存在，则表示有学生姓名重复，此时不能根据姓名确定身份证号码，予以删除
            stuCard_dict[v]=k
        else:
            del stuCard_dict[v]

    for pdfName in filenames:
        if os.path.isfile(work_dir / "pdf.png"):
            os.remove(work_dir / "pdf.png")
        pyMuPDF_fitz(work_dir / pdfName,work_dir / "pdf.png")#调用函数，将pdf的第一页转成图片
        model_root = Path(__file__).resolve().parents[2] / 'onnx' 
        ocr = CnOcr(rec_root=model_root / "cnocr",det_root=model_root / "cnstd")
        res = ocr.ocr(work_dir / "pdf.png")#识别图片
        pdfRename(work_dir,pdfName,res,studentInfo,stuNames_dict,stuCard_dict)  #重命名

    file_name = f"{schoolName}{datetime.now().strftime('%Y%m%d_%H%M%S')}{random.randint(10000,100000)}.zip"
    archive_path = UPLOAD_DIR / file_name
    shutil.make_archive(str(archive_path.with_suffix("")), "zip", work_dir)
    shutil.rmtree(work_dir)

    # 写 JSON
    data = json.loads(JSON_FILE.read_text(encoding="utf-8")) if JSON_FILE.exists() else {}
    data[code] = file_name
    JSON_FILE.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")


def repack_zip(zip_bytes: bytes, schoolName: str,studentInfo:str) -> str:
    """主线程：立刻返回 code，任务甩给后台"""
    code = f"{random.randint(1000, 9999)}"
    work_dir = UPLOAD_DIR / code
    work_dir.mkdir(exist_ok=True)

    tmp_zip = work_dir / "original.zip"
    tmp_zip.write_bytes(zip_bytes)

    executor.submit(_background_repack, zip_bytes, schoolName,studentInfo, code,tmp_zip)

    
    if not JSON_FILE.exists():
        JSON_FILE.write_text("{}", encoding="utf-8")    
    data = json.loads(JSON_FILE.read_text(encoding="utf-8"))
    data[code] = schoolName+"转换中,请稍候……"
    JSON_FILE.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")
    return code


@requestPdfRename.post("/pdfRenameUpload")
async def pdf_rename_upload(
    file: UploadFile = File(...),
    schoolName: str = Form(...),     # 前端同时传的学校名，先留着备用
    studentInfo: str = Form(...)
):
    if not file.filename.lower().endswith(".zip"):
        raise HTTPException(400, "只接受 zip 文件")

    content = await file.read()
    try:
        code = repack_zip(content,schoolName,studentInfo)
    except zipfile.BadZipFile:
        raise HTTPException(400, "非法 zip 包")
    return JSONResponse(content={"code": code})

@requestPdfRename.get("/haveTranslatedList")
def list_translated():
    translated_JsonFile = UPLOAD_DIR / "translated.json"
    if not translated_JsonFile.exists():
        translated_JsonFile.write_text("{}", encoding="utf-8")

    try:
        with translated_JsonFile.open("r", encoding="utf-8") as f:
            data = json.load(f)
    except json.JSONDecodeError as e:
        raise HTTPException(status_code=500, detail=f"JSON decode error: {e}")

    # 返回所有的值（假设 JSON 最外层是对象）
    return list(data.values())